/* @Author karl
** @Date 2025/11/1 10:58:00
** @Desc 限流器方法实现
 */
package g_phoenixLimiter

import (
	"fmt"
	"github.com/dgraph-io/ristretto/v2"
	"golang.org/x/time/rate"
	"sync"
	"sync/atomic"
	"time"
)

/*
* 	1. 默认按小时进行处理相关请求数据
*	2. 按小时进行清空相关缓存（范围不需要过于宽泛，导致缓存空间不够）
**/

type PhoenixLimiter interface {
	Check(key string, args ...interface{}) LimiterResult
	GetStats() RateLimiterStats
	SetConfig(config PhoenixLimiterConfig)
	GetBlacklist() []BlacklistEntity
	RemoveFromBlacklist(key string)
	Close()
}

type EnhancedPhoenixLimiter struct {
	config         *PhoenixLimiterConfig
	limiterCache   *ristretto.Cache[string, *rate.Limiter]
	blacklistCache *ristretto.Cache[string, BlacklistEntity]

	// 统计信息
	totalRequests    int64
	rejectedRequests int64
	blacklistedCount int32

	mu       sync.RWMutex
	stopChan chan struct{}
	once     sync.Once
}

func NewEnhancedPhoenixLimiter(config PhoenixLimiterConfig) *EnhancedPhoenixLimiter {
	// 设置默认值
	if config.BlacklistDuration == 0 {
		config.BlacklistDuration = 30 * time.Second // 改为30秒
	}
	if config.RecordDuration == 0 {
		config.RecordDuration = 30 * time.Second // 改为30秒
	}
	if config.PermitsPerSecond == 0 {
		config.PermitsPerSecond = 2 // 改为每秒2次
	}
	if config.BlackListCount == 0 {
		config.BlackListCount = 1 // 改为1次超限即拉黑
	}

	limiterCache, _ := ristretto.NewCache(&ristretto.Config[string, *rate.Limiter]{
		NumCounters: 1e7,
		MaxCost:     1 << 30,
		BufferItems: 64,
	})

	blacklistCache, _ := ristretto.NewCache(&ristretto.Config[string, BlacklistEntity]{
		NumCounters: 1e7,
		MaxCost:     1 << 30,
		BufferItems: 64,
	})

	limiter := &EnhancedPhoenixLimiter{
		config:         &config,
		limiterCache:   limiterCache,
		blacklistCache: blacklistCache,
		stopChan:       make(chan struct{}),
	}

	return limiter
}

func (epl *EnhancedPhoenixLimiter) Check(key string, args ...interface{}) LimiterResult {
	if !epl.config.Enable {
		return LimiterResult{Allowed: true, Reason: "", Message: "限流器未启动"}
	}

	atomic.AddInt64(&epl.totalRequests, 1)

	// 统一使用提取的key
	limitKey := key
	if epl.config.keyExtractor != nil {
		limitKey = epl.config.keyExtractor(args...)
	}
	if limitKey == "" {
		limitKey = "unknown"
	}

	// 1. 黑名单检查
	if epl.config.BlackListCount > 0 {
		if entity, exists := epl.blacklistCache.Get(limitKey); exists {
			if time.Now().Before(entity.ExpireTime) {
				if entity.Counts >= epl.config.BlackListCount {
					atomic.AddInt64(&epl.rejectedRequests, 1)
					return LimiterResult{
						Allowed: false,
						Reason:  "黑名单拦截",
						Message: fmt.Sprintf("IP %s 处于黑名单中，%.0f分钟后解除", limitKey, time.Until(entity.ExpireTime).Minutes()),
					}
				}
			} else {
				// 黑名单过期，移除
				epl.blacklistCache.Del(limitKey)
				atomic.AddInt32(&epl.blacklistedCount, -1)
			}
		}
	}

	// 2. 获取或创建限流器
	limiter, err := epl.getLimiter(limitKey)
	if err != nil {
		// 如果限流器创建失败，根据配置决定是否放行
		atomic.AddInt64(&epl.rejectedRequests, 1)
		return LimiterResult{
			Allowed: false,
			Reason:  "系统错误",
			Message: "限流器创建失败",
		}
	}

	// 3. 限流检查
	if !limiter.Allow() {
		atomic.AddInt64(&epl.rejectedRequests, 1)

		// 更新黑名单计数
		if epl.config.BlackListCount > 0 {
			epl.updateBlacklist(limitKey)
		}

		// 执行降级函数
		if epl.config.FallbackFunc != nil {
			result, err := epl.config.FallbackFunc(limitKey, "请求频率超限")
			if err == nil {
				return LimiterResult{
					Allowed: false,
					Reason:  "请求频率超限",
					Message: fmt.Sprintf("降级返回: %v", result),
				}
			}
		}

		return LimiterResult{
			Allowed: false,
			Reason:  "请求频率超限",
			Message: fmt.Sprintf("IP %s 请求过于频繁，请稍后再试", limitKey),
		}
	}

	return LimiterResult{
		Allowed: true,
		Reason:  "",
		Message: "请求允许通过",
	}
}

func (epl *EnhancedPhoenixLimiter) getLimiter(key string) (*rate.Limiter, error) {
	if limiter, exists := epl.limiterCache.Get(key); exists {
		return limiter, nil
	}

	limiter := rate.NewLimiter(rate.Limit(epl.config.PermitsPerSecond), epl.config.PermitsPerSecond)
	success := epl.limiterCache.SetWithTTL(key, limiter, 1, epl.config.RecordDuration)
	if !success {
		return nil, fmt.Errorf("failed to cache limiter")
	}

	return limiter, nil
}

func (epl *EnhancedPhoenixLimiter) updateBlacklist(key string) {
	epl.mu.Lock()
	defer epl.mu.Unlock()

	var entity BlacklistEntity
	if existing, exists := epl.blacklistCache.Get(key); exists {
		entity = existing
		entity.Counts++
	} else {
		entity = BlacklistEntity{
			Key:        key,
			Counts:     1,
			AddTime:    time.Now(),
			ExpireTime: time.Now().Add(epl.config.BlacklistDuration),
		}
		atomic.AddInt32(&epl.blacklistedCount, 1)
	}

	epl.blacklistCache.SetWithTTL(key, entity, 1, epl.config.BlacklistDuration)

	// 自动解除逻辑 - 5分钟后自动释放
	time.AfterFunc(epl.config.BlacklistDuration, func() {
		epl.blacklistCache.Del(key)
		atomic.AddInt32(&epl.blacklistedCount, -1)
	})
}

func (srl *EnhancedPhoenixLimiter) GetStats() RateLimiterStats {
	srl.mu.RLock()
	defer srl.mu.RUnlock()

	total := atomic.LoadInt64(&srl.totalRequests)
	rejected := atomic.LoadInt64(&srl.rejectedRequests)
	blacked := atomic.LoadInt32(&srl.blacklistedCount)
	rejectionRate := 0.0
	if total > 0 {
		rejectionRate = float64(rejected) / float64(total)
	}
	return RateLimiterStats{
		TotalRequests:    total,
		RejectedRequests: rejected,
		RejectionRate:    rejectionRate,
		BlacklistedCount: blacked,
		UpdateTime:       time.Now(),
	}
}
func (srl *EnhancedPhoenixLimiter) SetConfig(config PhoenixLimiterConfig) {
	srl.config = &config
}
func (srl *EnhancedPhoenixLimiter) GetBlacklist() []BlacklistEntity {
	// 下午再进行复杂的处理
	entries := make([]BlacklistEntity, 0)
	return entries
}
func (srl *EnhancedPhoenixLimiter) RemoveFromBlacklist(key string) {
	srl.mu.Lock()
	defer srl.mu.Unlock()
	srl.blacklistCache.Del(key)
	if srl.blacklistedCount > 0 {
		atomic.AddInt32(&srl.blacklistedCount, -1)
	}
	fmt.Printf("用户%s已经从黑名单中移除", key)
}
func (srl *EnhancedPhoenixLimiter) Close() {
	srl.once.Do(func() {
		close(srl.stopChan)
		srl.blacklistCache.Close()
		srl.limiterCache.Close()
		fmt.Println("限流器关闭")
	})
}
